from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from typing import List, Sequence
from langgraph.graph import END, MessageGraph
from llm import llm, llm0,llm0_1


# Generate
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "你是一名小说助理，任务是写出优秀的小说"
            "根据用户的要求，尽可能写出最好的小说"
            "如果用户提出了批评意见，请用你之前尝试过的修改版进行回复。",
        ),
        MessagesPlaceholder(variable_name="messages"),
    ]
)
generate = prompt | llm

# Reflect
reflection_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "你是一位为小说评分的老师。为用户提交的小说生成批语和建议"
            "提供详细的建议，包括长度、深度、风格等方面的要求",
        ),
        MessagesPlaceholder(variable_name="messages"),
    ]
)
reflect = reflection_prompt | llm


res = reflect.invoke({"messages": translated})
